# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import importlib
import logging
import typing

from absl import app, flags
from recordclass import dataobject

from coin.proto.coin_executor_pb2 import PassiveExecutorConfig
from coin.proto.coin_request_pb2 import StrategyRequestProto
from coin.strategy.executor.passive.factory import construct_mm_executor
from coin.experimental.jaewon.strat.strategy_driver import (StrategyBase,
                                                            StrategyDriver,
                                                            StrategyDriverConfig)
from coin.experimental.jaewon.strat.pass_model import PassModel, PassModelConfig


class PassModelStrategyConfig(dataobject):
  trade_product: typing.Any
  trade_sub_req: typing.Any
  executor_params: PassiveExecutorConfig
  model_config: PassModelConfig
  init_og_after_min: typing.Union[int, float, None] = None
  exit_after_min: typing.Union[int, float, None] = None
  og_params: typing.Any = None
  feed_checker_config: typing.Any = None
  strategy_request: StrategyRequestProto = None


class PassModelStrategy(StrategyBase):
  def __init__(self, config, dry_run=True, logger=None):
    super().__init__()

    self._logger = logger or logging.getLogger('PassModelStrategy')
    self._straegy_logger = None
    self._config = config

    self._model = PassModel(self._config.model_config,
                            pos_getter=self._get_pos,
                            trigger_callback=self.on_model_trigger)
    self._executor = None
    self._dry_run = dry_run
    self._exit_phase = False

  def prepare(self, feed_subsystem):
    self._model.prepare(feed_subsystem)

  def on_og_reset(self, order_gateway):
    self._logger.info('Intializing executor')
    self._executor = construct_mm_executor(self._config.trade_sub_req,
                                           self._config.executor_params,
                                           order_gateway,
                                           self._logger)
    self.strategy_logger.write_executor_config(self._config.executor_params)

  def on_model_trigger(self):
    if self._exit_phase or not self._executor:
      return

    if self._dry_run:
      self._logger.info('product: %s, sell: %s, buy: %s, sell_pull: %s, buy_pull: %s' %
                        (self._config.trade_product,
                         self._model.sell_price,
                         self._model.buy_price,
                         self._model.sell_price_pull,
                         self._model.buy_price_pull))

    self._executor.manage_mm_orders(sell_price=self._model.sell_price,
                                    buy_price=self._model.buy_price,
                                    sell_price_pull=self._model.sell_price_pull,
                                    buy_price_pull=self._model.buy_price_pull)

  def on_exit(self):
    self._exit_phase = True

  def _get_pos(self):
    if not self._executor:
      return float('nan')
    return self._executor.query_unbiased_position()

  @property
  def strategy_request(self):
    return self._config.strategy_request

  def get_driver_config(self, driver_config=None):
    driver_config = driver_config or StrategyDriverConfig()
    # TODO(jaewon): Support currency
    driver_config.currencies += [self._config.trade_product.base]
    driver_config.products += [self._config.trade_product]
    if driver_config.sub_req is not None:
      assert driver_config.sub_req == self._config.trade_sub_req
    driver_config.sub_req = self._config.trade_sub_req
    return driver_config


def main(argv):
  FLAGS = flags.FLAGS

  # TODO(jaewon): Make it json
  config_mod = importlib.import_module('coin.experimental.jaewon.strat.config.' + FLAGS.config_mod)
  strategy_config = config_mod.strategy_config()
  logging.info(str(strategy_config))

  init_og_after = datetime.timedelta(minutes=strategy_config.init_og_after_min)
  exit_after = datetime.timedelta(minutes=strategy_config.init_og_after_min
                                  + strategy_config.exit_after_min)

  strategy = PassModelStrategy(strategy_config, dry_run=FLAGS.dry_run)

  driver_config = strategy.get_driver_config()
  driver = StrategyDriver(strategy, driver_config)
  driver.register_signal_handler()

  return driver.run_live(strategy=driver,
                         on_order_gateway_reset_callback=driver.on_og_reset,
                         og_config_name=FLAGS.og_config_name,
                         og_params=driver_config.og_params,
                         init_order_gateway_after=init_og_after,
                         exit_after=exit_after,
                         feed_checker_config=strategy_config.feed_checker_config,
                         og_log_callback=driver.strategy_logger.write_og_log)


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')
  logging.getLogger('requests').setLevel(logging.WARNING)
  logging.getLogger('urllib3').setLevel(logging.WARNING)

  flags.DEFINE_string('og_config_name', None, 'OrderGateway config name')

  flags.DEFINE_string('config_mod', None, 'Config module')

  flags.DEFINE_boolean('dry_run', False, 'dry_run')

  app.run(main)
